// dllmain.cpp : Defines the entry point for the DLL application.
#include "stdafx.h"
#include "message.h"
#include <windows.h>
#include <stdio.h>
#include <process.h>
#include <list>
#include "..\include\detours.h"
using namespace std;

static CRITICAL_SECTION s_csPipe;                       // Guards access to hPipe.
static CRITICAL_SECTION s_csMessage;
static CRITICAL_SECTION s_csHeap;
static HANDLE           s_hPipe = INVALID_HANDLE_VALUE;
static DWORD            s_nProcessId = 0;

static HANDLE			s_heap = NULL;
static BOOL				s_endThread = FALSE;
static HANDLE			s_hThread = NULL;
struct MessageBuf{
	LogMessage buf[BUF_SIZE];
	DWORD messageCount;
};
//real allocator used to bypass hook.
static HANDLE InitHeap()
{
	if(s_heap == NULL){
		s_heap = HeapCreate(0,0,0);
	}
	return s_heap;
};
template <typename  T>
class real_allocator //: public allocator<T>
{  
public:  
    typedef T                 value_type;
    typedef value_type*       pointer;
    typedef const value_type* const_pointer;
    typedef value_type&       reference;
    typedef const value_type& const_reference;
    typedef std::size_t       size_type;
    typedef std::ptrdiff_t    difference_type;  
     
    real_allocator() {}  
    real_allocator(const real_allocator<T>&) _THROW0(){}
	template<class _Other>
	real_allocator(const real_allocator<_Other>&) _THROW0()
	{	// construct from a related allocator (do nothing)
	}
    ~real_allocator() {}  

    // rebind the class type u
    template <class u>
    struct rebind{
        typedef real_allocator<u>  other;
    };
     
    pointer address(reference x) const { return &x; }
    const_pointer address(const_reference x) const
    {   
        return &x;   
    }

    pointer allocate(size_type n, const_pointer = 0)
    {  
        if(s_heap == NULL) InitHeap();		
		EnterCriticalSection(&s_csHeap);
		//set flag to 0x50000163 to avoid recursive all HeapAlloc on windows 2008
		void* p = Real_HeapAlloc(s_heap, 0x50000163, n*sizeof(T));
		LeaveCriticalSection(&s_csHeap);
        return static_cast<pointer>(p);  
    }  
 
    void deallocate(pointer p, size_type)
    {  
        EnterCriticalSection(&s_csHeap);
		Real_HeapFree(s_heap, 0x50000063, p);
		LeaveCriticalSection(&s_csHeap);
    }  
 
    size_type max_size() const
    {   
        return static_cast<size_type>(-1) / sizeof(value_type);
        //return size_type(UNIT_MAX/sizeof(T));
    }  
     
    void construct(pointer p, const value_type& x)
    {
        new(p) value_type(x);   
    }  
 
    void destroy(pointer p) { p->~value_type(); }  

private:  
    void operator = (const real_allocator&);

}; 


template<typename _T1, typename _T2>
inline bool operator == (const real_allocator<_T1>&, const real_allocator<_T2>&)
{
    return true;
}

template<typename _T1, typename _T2>
inline bool operator != (const real_allocator<_T1>&, const real_allocator<_T2>&)
{
    return false;
}
static list<MessageBuf, real_allocator<MessageBuf>>* s_plogMessages = NULL; 
static BYTE dummyLogMessage[sizeof(*s_plogMessages)];

LPVOID (WINAPI *
        Real_HeapAlloc)(HANDLE hHeap, DWORD dwFlags, SIZE_T dwBytes)
    = HeapAlloc;
BOOL (WINAPI *
        Real_HeapFree) (HANDLE hHeap, DWORD dwFlags, LPVOID lpMem)
    = HeapFree;

LPVOID WINAPI Mine_HeapAlloc(HANDLE hHeap, DWORD dwFlags, SIZE_T dwBytes)
{
    LPVOID rv = 0;
	rv = Real_HeapAlloc(hHeap, dwFlags, dwBytes);
    __try {
		LogMessage message;
		ZeroMemory(&message, sizeof(message));
		LogMessage::Stackinfo *pMS = &message.callStack;
		DWORD asmBP;
		DWORD TEB;
		_asm {
			mov asmBP, ebp;
			mov eax,dword ptr fs:[18h];
			mov TEB, eax;
		}
		HANDLE hProcess = GetCurrentProcess();
		for(int i=0; i<MEMORY_STACKS; i++) {
			DWORD codePointer;
			codePointer = ((DWORD *)asmBP)[1];
			pMS->stack[i] = codePointer;
			DWORD newBP;
			newBP = *(DWORD *)(asmBP);
			if ( newBP <= asmBP || newBP > ((DWORD *)TEB)[1]){
				break;
			}
			asmBP = newBP;
		}
        
		message.function = FunHeapAlloc;
		message.memAlloc = rv;
		message.hHeap = hHeap;
		message.size = dwBytes;
		message.nProcessId = GetCurrentProcessId();
		//save message
		EnterCriticalSection(&s_csMessage);		
		s_plogMessages->back().buf[s_plogMessages->back().messageCount] = message;
		s_plogMessages->back().messageCount++;
		if (s_plogMessages->back().messageCount >=BUF_SIZE){
			s_plogMessages->push_back(MessageBuf());
		}
		LeaveCriticalSection(&s_csMessage);
    } __finally {

    };
    return rv;
}
BOOL WINAPI Mine_HeapFree(HANDLE hHeap, DWORD dwFlags, LPVOID lpMem)
{
	BOOL rv = FALSE;
	rv = Real_HeapFree(hHeap, dwFlags, lpMem);
	__try{
		LogMessage message;
		ZeroMemory(&message, sizeof(message));		
		message.function = FunHeapFree;
		message.memAlloc = lpMem;
		message.hHeap = hHeap;
		message.nProcessId = GetCurrentProcessId();
		//save message
		EnterCriticalSection(&s_csMessage);	
		s_plogMessages->back().buf[s_plogMessages->back().messageCount] = message;
		s_plogMessages->back().messageCount++;
		if (s_plogMessages->back().messageCount >=BUF_SIZE){
			s_plogMessages->push_back(MessageBuf());
		}
		LeaveCriticalSection(&s_csMessage);
		
	} __finally {

    };
	return rv;
}
static unsigned int WINAPI WorkerThread(LPVOID p)
{
	while (1){
		Sleep(5000);
		if (s_endThread){
			EnterCriticalSection(&s_csMessage);
			while (s_plogMessages->size() > 0){
					DWORD nbytes;
					WriteFile(s_hPipe,s_plogMessages->front().buf, 
						s_plogMessages->front().messageCount*sizeof(LogMessage),
						&nbytes,
						NULL);
					s_plogMessages->pop_front();
				}
			LeaveCriticalSection(&s_csMessage);
			break;
		}else{
			EnterCriticalSection(&s_csPipe);
			if(s_hPipe != INVALID_HANDLE_VALUE){
				EnterCriticalSection(&s_csMessage);
				if(s_plogMessages->back().messageCount > 0){
					s_plogMessages->push_back(MessageBuf());
				}
				LeaveCriticalSection(&s_csMessage);
				while (s_plogMessages->size() > 1){
					DWORD nbytes;
					WriteFile(s_hPipe,s_plogMessages->front().buf, 
						s_plogMessages->front().messageCount*sizeof(LogMessage),
						&nbytes,
						NULL);
					s_plogMessages->pop_front();
				}
			}
			LeaveCriticalSection(&s_csPipe);
		}
	}
	return 0;
}
LONG AttachDetours(VOID)
{
    DetourTransactionBegin();
    DetourUpdateThread(GetCurrentThread());

    DetourAttach(&(PVOID&)Real_HeapAlloc, Mine_HeapAlloc);
	DetourAttach(&(PVOID&)Real_HeapFree, Mine_HeapFree);
    return DetourTransactionCommit();
}

LONG DetachDetours(VOID)
{
    DetourTransactionBegin();
    DetourUpdateThread(GetCurrentThread());

    DetourDetach(&(PVOID&)Real_HeapAlloc, Mine_HeapAlloc);
	DetourDetach(&(PVOID&)Real_HeapFree, Mine_HeapFree);
    return DetourTransactionCommit();
}
BOOL APIENTRY DllMain( HMODULE hModule,
                       DWORD  ul_reason_for_call,
                       LPVOID lpReserved
					 )
{
	switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:
		OutputDebugStringA("Attach process\n");
		InitHeap();
		InitializeCriticalSection(&s_csHeap);
		InitializeCriticalSection(&s_csPipe);
		InitializeCriticalSection(&s_csMessage);
		s_plogMessages = (list<MessageBuf, real_allocator<MessageBuf>>*)dummyLogMessage;
		new (s_plogMessages) list<MessageBuf, real_allocator<MessageBuf>>();
		s_plogMessages->push_back(MessageBuf());

		if(s_hPipe == INVALID_HANDLE_VALUE){
			s_hPipe = CreateFile(TRACE_PIPE_NAME,
				GENERIC_WRITE, 0, NULL, OPEN_EXISTING, 0, NULL);
			if(s_hPipe != INVALID_HANDLE_VALUE){
				DWORD dwMode = PIPE_READMODE_MESSAGE;
				SetNamedPipeHandleState(s_hPipe, &dwMode, NULL, NULL);
			}
		}
		s_hThread = (HANDLE)_beginthreadex(NULL,0,WorkerThread,NULL,0,NULL);
		AttachDetours();
		break;
	case DLL_THREAD_ATTACH:
		break;
	case DLL_THREAD_DETACH:
		break;
	case DLL_PROCESS_DETACH:
		DetachDetours();
		OutputDebugStringA("detach process\n");
		s_endThread = TRUE;
		WaitForSingleObject(s_hThread,INFINITE);
		FlushFileBuffers(s_hPipe);
		CloseHandle(s_hPipe);
        s_hPipe = INVALID_HANDLE_VALUE;
		break;
	}
	return TRUE;
}

